package main

import (
	"bytes"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"
)

func init() {
	sql.Register("bigqueryproxy", &BigQueryProxyDriver{})
}

// BigQueryProxyDriver 实现 driver.Driver 接口
type BigQueryProxyDriver struct{}

// Open 打开一个新的连接
func (d *BigQueryProxyDriver) Open(dsn string) (driver.Conn, error) {
	connName, serverURL := parseBigQueryDSN(dsn)

	config := ConnConfig{
		ServerURL:  serverURL,
		Timeout:    time.Second * 30,
		DBType:     "bigquery",
		Connection: connName,
	}

	return &BigQueryConn{
		config: config,
		client: &http.Client{
			Timeout: config.Timeout,
		},
	}, nil
}

// 解析BigQuery DSN
// 格式: [connection_name@]host:port
func parseBigQueryDSN(dsn string) (connName, serverURL string) {
	connName = "default"

	// 检查是否指定了连接名称
	parts := splitAtSign(dsn)
	if len(parts) > 1 {
		connName = parts[0]
		dsn = parts[1]
	}

	// 确保URL包含协议
	if !hasProtocol(dsn) {
		dsn = "http://" + dsn
	}

	return connName, dsn
}

// BigQueryConn 实现了driver.Conn接口
type BigQueryConn struct {
	config ConnConfig
	client *http.Client
}

// sendRequest 发送请求到服务器
func (c *BigQueryConn) sendRequest(req SQLRequest) (*SQLResponse, error) {
	// 设置数据库类型和连接名称
	req.DBType = c.config.DBType
	req.Connection = c.config.Connection

	data, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequest("POST", buildRequestURL(c.config.ServerURL), bytes.NewBuffer(data))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("服务器错误: %s - %s", resp.Status, string(body))
	}

	var sqlResp SQLResponse
	if err := json.NewDecoder(resp.Body).Decode(&sqlResp); err != nil {
		return nil, err
	}

	if sqlResp.Error != "" {
		return nil, errors.New(sqlResp.Error)
	}

	return &sqlResp, nil
}

// Prepare 实现 driver.Conn 接口的 Prepare 方法
func (c *BigQueryConn) Prepare(query string) (driver.Stmt, error) {
	req := SQLRequest{
		SQL:       query,
		Operation: "prepare",
	}

	resp, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	if len(resp.Rows) == 0 || len(resp.Rows[0]) == 0 {
		return nil, errors.New("准备语句失败: 无有效ID返回")
	}

	stmtID, ok := resp.Rows[0][0].(string)
	if !ok {
		return nil, errors.New("准备语句失败: 返回的ID格式无效")
	}

	return &BigQueryStmt{
		conn:   c,
		query:  query,
		stmtID: stmtID,
	}, nil
}

// Close 实现 driver.Conn 接口的 Close 方法
func (c *BigQueryConn) Close() error {
	return nil
}

// Begin 实现 driver.Conn 接口的 Begin 方法
func (c *BigQueryConn) Begin() (driver.Tx, error) {
	return nil, errors.New("BigQuery不支持事务")
}

// Exec 实现扩展接口，直接执行SQL语句
func (c *BigQueryConn) Exec(query string, args []driver.Value) (driver.Result, error) {
	dargs := valuesToInterface(args)
	req := SQLRequest{
		SQL:       query,
		Args:      dargs,
		Operation: "exec",
	}

	resp, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	if len(resp.Rows) == 0 || len(resp.Rows[0]) < 2 {
		return nil, errors.New("执行失败: 响应格式无效")
	}

	// 获取LastInsertId和RowsAffected
	var rowsAffected int64

	if v, ok := resp.Rows[0][1].(float64); ok {
		rowsAffected = int64(v)
	}

	return &BigQueryResult{
		rowsAffected: rowsAffected,
	}, nil
}

// Query 实现扩展接口，直接执行查询
func (c *BigQueryConn) Query(query string, args []driver.Value) (driver.Rows, error) {
	dargs := valuesToInterface(args)
	req := SQLRequest{
		SQL:       query,
		Args:      dargs,
		Operation: "query",
	}

	resp, err := c.sendRequest(req)
	if err != nil {
		return nil, err
	}

	return &BigQueryRows{
		columns: resp.Columns,
		rows:    resp.Rows,
		rowIdx:  -1,
	}, nil
}

// BigQueryResult 实现 driver.Result 接口
type BigQueryResult struct {
	rowsAffected int64
}

// LastInsertId BigQuery不支持获取最后插入的ID
func (r *BigQueryResult) LastInsertId() (int64, error) {
	return 0, errors.New("BigQuery不支持LastInsertId")
}

// RowsAffected 返回受影响的行数
func (r *BigQueryResult) RowsAffected() (int64, error) {
	return r.rowsAffected, nil
}

// BigQueryStmt 实现 driver.Stmt 接口
type BigQueryStmt struct {
	conn   *BigQueryConn
	query  string
	stmtID string
}

// Close 关闭语句
func (s *BigQueryStmt) Close() error {
	return nil
}

// NumInput 返回参数数量
func (s *BigQueryStmt) NumInput() int {
	return -1
}

// Exec 执行预处理语句
func (s *BigQueryStmt) Exec(args []driver.Value) (driver.Result, error) {
	dargs := valuesToInterface(args)
	req := SQLRequest{
		SQL:       s.query,
		Args:      dargs,
		Operation: "exec",
	}

	resp, err := s.conn.sendRequest(req)
	if err != nil {
		return nil, err
	}

	if len(resp.Rows) == 0 || len(resp.Rows[0]) < 2 {
		return nil, errors.New("执行失败: 响应格式无效")
	}

	// 获取RowsAffected
	var rowsAffected int64

	if v, ok := resp.Rows[0][1].(float64); ok {
		rowsAffected = int64(v)
	}

	return &BigQueryResult{
		rowsAffected: rowsAffected,
	}, nil
}

// Query 执行预处理查询
func (s *BigQueryStmt) Query(args []driver.Value) (driver.Rows, error) {
	dargs := valuesToInterface(args)
	req := SQLRequest{
		SQL:       s.query,
		Args:      dargs,
		Operation: "query",
	}

	resp, err := s.conn.sendRequest(req)
	if err != nil {
		return nil, err
	}

	return &BigQueryRows{
		columns: resp.Columns,
		rows:    resp.Rows,
		rowIdx:  -1,
	}, nil
}

// BigQueryRows 实现 driver.Rows 接口
type BigQueryRows struct {
	columns []string
	rows    [][]interface{}
	rowIdx  int
}

// Columns 返回列名
func (r *BigQueryRows) Columns() []string {
	return r.columns
}

// Close 关闭结果集
func (r *BigQueryRows) Close() error {
	return nil
}

// Next 获取下一行数据
func (r *BigQueryRows) Next(dest []driver.Value) error {
	r.rowIdx++
	if r.rowIdx >= len(r.rows) {
		return io.EOF
	}

	row := r.rows[r.rowIdx]
	if len(row) != len(dest) {
		return fmt.Errorf("列数不匹配: 需要 %d, 返回 %d", len(dest), len(row))
	}

	for i, v := range row {
		dest[i] = driver.Value(v)
	}

	return nil
} 